"""LLM Plays Pokemon 知识库基础模块"""

import os
import json
import time
import logging
import shutil
from typing import Dict, Any, Optional, List

logger = logging.getLogger(__name__)


class KnowledgeBase:
    """知识库类，用于存储和管理AI的长期记忆"""
    
    def __init__(self, storage_path: str, backup_dir: str, max_backups: int = 10):
        """初始化知识库
        
        Args:
            storage_path: 知识库存储路径
            backup_dir: 备份目录
            max_backups: 最大备份数量
        """
        self.storage_path = storage_path
        self.backup_dir = backup_dir
        self.max_backups = max_backups
        self.data = {}
        self.modified = False
        
        # 加载知识库
        self._load()
        
        logger.info("知识库初始化完成")
    
    def _load(self):
        """加载知识库"""
        if os.path.exists(self.storage_path):
            try:
                with open(self.storage_path, 'r', encoding='utf-8') as f:
                    self.data = json.load(f)
                logger.info(f"已加载知识库: {self.storage_path}")
            except Exception as e:
                logger.error(f"加载知识库失败: {e}")
                self.data = {}
        else:
            logger.info(f"知识库文件不存在: {self.storage_path}，创建新知识库")
            self.data = {}
            self.modified = True
    
    def save(self):
        """保存知识库"""
        if not self.modified:
            return
        
        try:
            # 确保目录存在
            os.makedirs(os.path.dirname(os.path.abspath(self.storage_path)), exist_ok=True)
            
            # 创建备份
            self._create_backup()
            
            # 保存知识库
            with open(self.storage_path, 'w', encoding='utf-8') as f:
                json.dump(self.data, f, indent=2, ensure_ascii=False)
            
            self.modified = False
            logger.info(f"已保存知识库: {self.storage_path}")
        except Exception as e:
            logger.error(f"保存知识库失败: {e}")
    
    def _create_backup(self):
        """创建知识库备份"""
        if not os.path.exists(self.storage_path):
            return
        
        try:
            # 确保备份目录存在
            os.makedirs(self.backup_dir, exist_ok=True)
            
            # 创建备份文件名
            timestamp = time.strftime("%Y%m%d_%H%M%S")
            backup_filename = f"knowledge_base_{timestamp}.json"
            backup_path = os.path.join(self.backup_dir, backup_filename)
            
            # 复制文件
            shutil.copy2(self.storage_path, backup_path)
            
            # 清理旧备份
            self._cleanup_old_backups()
            
            logger.info(f"已创建知识库备份: {backup_path}")
        except Exception as e:
            logger.error(f"创建知识库备份失败: {e}")
    
    def _cleanup_old_backups(self):
        """清理旧备份"""
        try:
            # 获取所有备份文件
            backup_files = [f for f in os.listdir(self.backup_dir) 
                           if f.startswith("knowledge_base_") and f.endswith(".json")]
            
            # 按修改时间排序
            backup_files.sort(key=lambda f: os.path.getmtime(os.path.join(self.backup_dir, f)), reverse=True)
            
            # 删除多余的备份
            for old_file in backup_files[self.max_backups:]:
                os.remove(os.path.join(self.backup_dir, old_file))
                logger.debug(f"已删除旧备份: {old_file}")
        except Exception as e:
            logger.error(f"清理旧备份失败: {e}")
    
    def get_section(self, section_id: str) -> Optional[str]:
        """获取知识库中的特定部分
        
        Args:
            section_id: 部分ID
            
        Returns:
            部分内容，如果不存在则返回None
        """
        return self.data.get(section_id)
    
    def set_section(self, section_id: str, content: str):
        """设置知识库中的特定部分
        
        Args:
            section_id: 部分ID
            content: 部分内容
        """
        self.data[section_id] = content
        self.modified = True
        logger.debug(f"已更新知识库部分: {section_id}")
    
    def delete_section(self, section_id: str) -> bool:
        """删除知识库中的特定部分
        
        Args:
            section_id: 部分ID
            
        Returns:
            是否成功删除
        """
        if section_id in self.data:
            del self.data[section_id]
            self.modified = True
            logger.debug(f"已删除知识库部分: {section_id}")
            return True
        return False
    
    def list_sections(self) -> List[str]:
        """列出所有部分ID
        
        Returns:
            部分ID列表
        """
        return list(self.data.keys())
    
    def get_all(self) -> Dict[str, str]:
        """获取整个知识库
        
        Returns:
            知识库数据
        """
        return self.data.copy()
    
    def clear(self):
        """清空知识库"""
        self.data = {}
        self.modified = True
        logger.info("已清空知识库")